import os
import numpy as np
import scipy.io.wavfile as wavfile
import tensorflow as tf
import shutil

# =================== config for spectrogram =========================
wav_dir = '/home/dai/Project/emotions/iemocap_4emo_wav'
spectrogram_dir = '/home/dai/Project/emotions/iemocap_4emo_spec'

# time of a frame
# frame_time = 0.04
count = 0


# =================== convert wav to spectrogram =====================
def cal_spectrogram(signals, frame_size, step_frac):
    """signals: tf.placeholder [batch_size, signal_length]"""
    # [batch_size, signal_length]
    # signals = tf.placeholder(tf.float32, shape=[None, None], name='signals')
    frame_step = int(frame_size * step_frac)
    stfts = tf.contrib.signal.stft(signals, frame_length=frame_size, frame_step=frame_step,
                                   fft_length=1600,
                                   window_fn=tf.contrib.signal.hamming_window, pad_end=True)
    magnitude_spectrogram = tf.abs(stfts)
    log_offset = 1e-6
    log_spectrogram = 10 * tf.log(magnitude_spectrogram + log_offset)
    return log_spectrogram

#
# def cal_spectrograms2(wavs, session_name, frame_time=0.04, step_frac=0.25):
#     global count
#     seq_lens = list()
#     signals = tf.placeholder(tf.float32, [1, None])
#     with tf.Session() as sess:
#         for wav in wavs:
#             rate, data = wavfile.read(wav)
#             data = data.reshape([1, -1])
#             frame_size = int(rate * frame_time)
#


def cal_spectrograms(wavs, frame_time=0.04, step_frac=0.25):
    global count
    spectrograms = list()
    seq_lens = list()
    signals = tf.placeholder(tf.float32, [1, None])
    # log_spectrogram = cal_spectrogram(signals, frame_size, step_frac)
    with tf.Session() as sess:
        for wav in wavs:
            rate, data = wavfile.read(wav)
            data = data.reshape([1, -1])
            frame_size = int(rate * frame_time)
            l_spectr = cal_spectrogram(signals, frame_size, step_frac)
            spectrogram = l_spectr.eval(feed_dict={signals: data}, session=sess)[0]
            time_scale, freq_scale = spectrogram.shape
            spectrograms.append(spectrogram)
            seq_lens.append(time_scale)
    return spectrograms, seq_lens


# def cal_spectrograms_dir(wav_dir):
#     wav_files = os.listdir(wav_dir)
#     wav_paths = [os.path.join(wav_dir, wav_file) for wav_file in wav_files if '.wav' in wav_file]
#     spectrograms, seq_lens = cal_spectrograms(wav_paths)
#     return spectrograms, seq_lens
#

def cal_session(session_name):
    setence_label_filename = os.path.join(wav_dir, session_name + '_sentence_label')
    wav_paths = list()
    labels = list()
    with open(setence_label_filename, 'r') as s_f:
        for line in s_f:
            eles = line.split()
            if len(eles) == 2:
                wav_path = os.path.join(wav_dir, session_name, eles[0]+'.wav')
                wav_paths.append(wav_path)
                labels.append(eles[1])
    spectrograms, seq_lens = cal_spectrograms(wav_paths)
    return spectrograms, seq_lens, labels


def cal_save_sessions():
    if os.path.exists(spectrogram_dir):
        shutil.rmtree(spectrogram_dir)
    os.makedirs(spectrogram_dir)
    session_names = ['Session1', 'Session2', 'Session3', 'Session4', 'Session5']

    # every elements in spectrograms_list is spectrograms of a session
    spectrograms_list = list()

    # seq_lens_list = list()
    max_seq_len = 0
    for session_name in session_names:
        spectrograms, seq_lens, labels = cal_session(session_name)
        max_seq_len = max_seq_len if max_seq_len > max(seq_lens) else max(seq_lens)
        seq_lens_f = os.path.join(spectrogram_dir, session_name + '_seq_len.npy')
        np.save(seq_lens_f, seq_lens)
        labels_f = os.path.join(spectrogram_dir, session_name + '_label.npy')
        np.save(labels_f, labels)
        spectrograms_list.append(spectrograms)
    for session_name, spectrograms in zip(session_names, spectrograms_list):
        spectrogram_pad_list = list()
        for spectrogram in spectrograms:
            time_scale, freq_scale = spectrogram.shape
            spectrogram_pad = np.zeros([max_seq_len, freq_scale])
            spectrogram_pad[:time_scale, :freq_scale] = spectrogram
            spectrogram_pad_list.append(spectrogram_pad)
        spectrogram_pad_np = np.array(spectrogram_pad_list)
        spectrogram_pad_np_f = os.path.join(spectrogram_dir, session_name + '_spectrogram.npy')
        np.save(spectrogram_pad_np_f, spectrogram_pad_np)


def main():
    # session_name = 'Session1'
    # spectrograms, seq_lens, labels = cal_session(session_name)
    # for spectrogram, seq_len, label in zip(spectrograms, seq_lens, labels):
    #     print(spectrogram.shape)
    #     print(seq_len)
    #     print(label)
    #     print()
    cal_save_sessions()


if __name__ == '__main__':
    main()






